from pygrib import open as pyopen
import os
from datetime import datetime, timedelta
import time
import numpy as np
from utils.config import apcp_path, input_path
import glob
from multiprocessing import Pool, cpu_count

# 组合反射率区域: latmin, latmax, lonmin, lonmax = 19, 28, 102, 114

# apcp全区域:  18.21, 27.8, 102.2, 113.79
# 实际选取区域: 27.8, 19.01, 102.2, 113.79
latmax, latmin, lonmin, lonmax = 27.8, 19, 102.2, 113.8


def get_grib2(file):
    grbs = pyopen(file)
    grbs.seek(0)
    try:
        grb = grbs.select(parameterName="Total precipitation")[0]
        area_data, lats, lons = grb.data(lat1=latmin, lat2=latmax, lon1=lonmin, lon2=lonmax)
        if area_data.shape == (880, 1160):
            return area_data
        else:
            print("bad apcp data:", file)
            return None
    except ValueError:
        print("BAD GRIB2 DATA", file)
        return None


def get_APCP(file, save=True):
    fdir, fname = os.path.split(file)

    stime = datetime.strptime(fname[12:24], '%Y%m%d%H%M')
    fmins = int(fname[-8:-5])
    if fmins > 0:  # 排除0场数据
        ftime = (stime + timedelta(minutes=fmins)).strftime('%Y%m%d%H%M')
        # lasttime = (stime + timedelta(minutes=(fmins-12))).strftime('%Y%m%d%H%M')
        lastfile = f"{fdir}/{fname[:-8]}{str(fmins - 12).rjust(3, '0')}.grb2"
        if os.path.exists(lastfile):
            try:
                fdata = get_grib2(file)

                ldata = get_grib2(lastfile)

                if fdata is not None and ldata is not None:
                    apcp_data = fdata - ldata
                    apcp_data[apcp_data < 0] = 0

                    sdir = f"{input_path}/{fname[:-10]}"
                    if not os.path.exists(sdir):
                        os.mkdir(sdir)
                    if save:
                        apcp_file = f"{sdir}/{ftime}.000"
                        valcnt = np.sum(apcp_data > 0)
                        valth = apcp_data.shape[0] * apcp_data.shape[1] * 0.2
                        if valcnt > valth:
                            print("save ", fmins, apcp_file)
                            np.savetxt(apcp_file, apcp_data, fmt='%.02f')
                    else:
                        return apcp_data

            except:
                print("Wrong", file)
                return None

        else:
            return None


if __name__ == '__main__':

    dirs = os.listdir(apcp_path)
    dirs.sort()
    for dir in dirs:
        if int(dir.split('GX-R1KM_')[1][:10]) < 2021070100:
            continue
        files = glob.glob(f"{apcp_path}/{dir}/*.grb2")
        files.sort()
        for f in files:
            get_APCP(f)

        # pool = Pool(processes=10)
        # pool.map(get_APCP, files)
        # pool.close()
        # pool.join()
